# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
from typing import Union, List
import random
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig

@GENERATOR_REGISTRY.register("ascend_generate_aclnn_moe_token_unpermute")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.tokens_num = 1
        self.top_k = 1
        self.hidden_size = 1
        self.case_dtype = None

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]],
            ) -> Union[InputCaseConfig, List[InputCaseConfig]]:
        if index == 2:
            self.tokens_num = input_case.shape[0]
            self.top_k = input_case.shape[1]
            if self.top_k>=512:
                top_k_list=[ 1, 4, 5, 8, 10, 12, 15, 19, 22, 186, 243, 357, 412 ]
                self.top_k=random.choice(top_k_list)
            if self.tokens_num==4096:
                self.top_k=8
            input_case.dtype = self.case_dtype
        if index == 0:
            self.hidden_size = input_case.shape[1]
            self.case_dtype = input_case.dtype
        return input_case
    
    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        case_config.inputs[0].shape = [self.tokens_num * self.top_k, self.hidden_size]
        case_config.inputs[1].shape = [self.tokens_num * self.top_k]
        case_config.inputs[2].shape = [self.tokens_num, self.top_k]
        # 值域[-100, 100]存在计算溢出，适当缩小值域范围，改为正确性计算值域范围 [-5, 5]
        case_config.inputs[0].range_values = [-5, 5]
        case_config.inputs[2].range_values = [-5, 5]
        return case_config
